fix: shape the value buffer for coordinate selections in sharded writes - #4284
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4284 +/- ##
=======================================
Coverage 94.21% 94.21%
=======================================
Files 92 92
Lines 12863 12871 +8
=======================================
+ Hits 12119 12127 +8
Misses 744 744
🚀 New features to boost your workflow:
|
|
🤖 AI text below 🤖 Code reviewFound 1 issue:
Missing-fix location: zarr-python/src/zarr/codecs/sharding.py Lines 798 to 806 in 01aac1f The reshape added on the async side: zarr-python/src/zarr/codecs/sharding.py Lines 1361 to 1374 in 01aac1f Sync decode twin already carrying the mirrored reshape: zarr-python/src/zarr/codecs/sharding.py Lines 1306 to 1309 in 01aac1f 🤖 Generated with Claude Code - If this code review was useful, please react with 👍. Otherwise, react with 👎. |
_encode_partial_sync derives its indexer the same way as _encode_partial_single and so hits the same coordinate-selection shape mismatch under FusedCodecPipeline. The regression test is parametrized over both pipelines.
|
Good catch, that was a real gap. Same reshape applied there, and the regression test is now parametrized over both pipelines. All four partial paths handle I confirmed the two Fused variants fail on the previous commit and pass on this one, so the parametrization is doing work rather than duplicating a passing case.
|
|
We hit this issue in icechunk CI too, and are currently sidestepping it There is also an extra fix in |
|
Thanks for the corroboration — useful to know it reproduces in icechunk CI and not just in a constructed case. The stable-sort change in |
|
this looks good, i'm going to approve + merge. there's a remaining bugged case that I will fix in a follow-up PR: import numpy as np, zarr
from zarr.storage import MemoryStore
from zarr.codecs import ShardingCodec, BytesCodec
a = zarr.create_array(
MemoryStore(), shape=(4, 4, 4), chunks=(2, 4, 4), dtype="int32",
serializer=ShardingCodec(chunk_shape=(2, 2, 2), codecs=(BytesCodec(),)),
compressors=None, fill_value=0,
)
a[:] = np.arange(64, dtype="int32").reshape(4, 4, 4)
a.oindex[np.array([3, 1, 2]), 1, np.array([0, 2])] = np.arange(6, dtype="int32").reshape(3, 2)
# ValueError: shape mismatch: value array of shape (1,2) could not be broadcast to indexing result of shape (1,) |
|
Thanks for confirming the bug hits icechunk too, that's useful beyond our regression case. Agreed the stable-sort change in |
…in partial writes The guard added in zarr-developers#4284 only reshaped the value when its shape equalled the re-derived CoordinateIndexer's sel_shape. An orthogonal selection that mixes an integer index with two or more array indices defeats that: OrthogonalIndexer drops the integer axis from the value but np.ix_ keeps it as a length-1 axis in the chunk selection, so the shapes differ in rank while agreeing in element count, the reshape was skipped, and the write still raised the shape-mismatch ValueError. The invariant is that a coordinate indexer addresses the value flat, so ravel any multi-dimensional value instead. Both partial-encode paths now share one helper for deriving the shard indexer and shaping the value, and the check is an isinstance on CoordinateIndexer so mypy types sel_shape. The regression test is parametrized over selections with an integer axis in each position, three array axes, and an unsorted selection spanning two shards. Closes zarr-developers#4315 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…in partial writes (#4316) * fix(sharding): flatten any multi-dim value for coordinate selections in partial writes The guard added in #4284 only reshaped the value when its shape equalled the re-derived CoordinateIndexer's sel_shape. An orthogonal selection that mixes an integer index with two or more array indices defeats that: OrthogonalIndexer drops the integer axis from the value but np.ix_ keeps it as a length-1 axis in the chunk selection, so the shapes differ in rank while agreeing in element count, the reshape was skipped, and the write still raised the shape-mismatch ValueError. The invariant is that a coordinate indexer addresses the value flat, so ravel any multi-dimensional value instead. Both partial-encode paths now share one helper for deriving the shard indexer and shaping the value, and the check is an isinstance on CoordinateIndexer so mypy types sel_shape. The regression test is parametrized over selections with an integer axis in each position, three array axes, and an unsorted selection spanning two shards. Closes #4315 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * docs: add changelog entry for #320 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * Rename 320.bugfix.md to 4316.bugfix.md * fix(sharding): ravel only a value shaped like the selection minus its unit axes Ravelling every multi-dimensional value for a coordinate selection was too lenient. A mask write with a (2, 2) value for four selected elements, or an orthogonal write with a spurious trailing axis, raises on an unsharded array but was silently accepted on a sharded one, because the element count matched and the shard-level selection cannot tell orthogonal from mask indexing. The value shape can. An np.ix_ selection has an N-D sel_shape and the caller's value is that shape minus the integer-indexed axes, which np.ix_ keeps as length-1 axes. Ravel exactly that shape and leave any other rank alone, so an invalid write fails the same way it does without sharding. Adds an error test for both leniencies and a positive case with a length-1 array axis next to an integer axis. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test: exercise sharded writes and bare integer axes in the indexing property tests The property tests could not have found the sharded orthogonal-write bugs: - test_oindex, test_mask_indexing and test_block_indexing skipped their set half on sharded arrays with assume(zarray.shards is None), added in #2825 when the bug was first seen and never lifted. test_vindex had its set half commented out. - orthogonal_indices wrapped every bare integer as a one-element array, so zarr never received an integer index and OrthogonalIndexer's dropped-axis path was unreachable. basic_indices(min_dims=1) never yields an integer either, so that branch was dead. - arrays() only drew a shard shape when every axis had a chunk strictly between 1 and the axis length, on top of the v3 and regular-grid draws: 2 of 500 test_oindex examples were sharded. Lift the skips, draw integers explicitly and give the numpy indexer the same dropped-axis result, enable the vindex write with a duplicate-point filter, and let any chunk that fits the array be sharded (33 of 500 now). With these changes test_oindex fails against the code before this PR with the mixed-integer shape mismatch, and passes with it. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test: assert invalid value ranks are rejected on chunked and sharded arrays alike Replace the sharded-only wrong-rank test in test_sharding.py and the GH2469 one-off in test_indexing.py with one parametrized error test: a coordinate write with twice the elements, a mask write with a 2-D value, and an orthogonal write with an extra axis each raise ValueError on chunked and sharded arrays under both codec pipelines. The property under test is that storage layout does not change which writes are rejected, which a sharded-only test could not state. zarr_array_from_numpy_array grows a shards argument for it. Only the rejection is asserted; a write that fails inside the chunk merge may already have touched other chunks on a chunked array. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Summary
Nightly
Slow Hypothesis CIon main filed #4280 on 2026-08-22 and hit it again on 2026-08-25 (run 32792180919):ValueError: shape mismatch: value array of shape (3,1) could not be broadcast to indexing result of shape (3,). An orthogonal set on a sharded array with two array-indexed dimensions reproduces it:OrthogonalIndexerhands such a chunk selection down as annp.ix_pair (indexing.py:989).get_indexerreads it back as a coordinate selection, whose projections addressshard_arrayflat while the caller shaped it likesel_shape._decode_partial_singlereshapesouttosel_shapeon the way out (sharding.py:1085); this does the same on the way in. Reads were unaffected.For reviewers
Guarded on
shard_array.shape == sel_shape, so a flat value passes through. Checked against a numpynp.ix_oracle over 1568 combinations of chunk grid, sharding nesting, per-dimension selector: 96 failures on main, 0 after, every one a write with two array-indexed dimensions.Author attestation
TODO
changes/